import pandas as pd
import peewee as pw
import os
from datetime import datetime

from models import create_tables, Transaction

# 数据库连接配置
db = pw.MySQLDatabase(
    'qkl',
    user='root',
    password='1473992017',
    host='localhost',
    port=3306
)


def import_data_from_csv(csv_path):
    """从CSV文件导入数据到数据库"""
    try:
        # 检查文件是否存在
        if not os.path.exists(csv_path):
            print(f"错误: 文件 {csv_path} 不存在")
            return False
        
        # 读取CSV文件
        print(f"正在读取CSV文件: {csv_path}")
        df = pd.read_csv(csv_path)
        
        # 检查必要的列是否存在
        required_columns = ['TxHash', 'BlockHeight', 'TimeStamp', 'From', 'To', 'Value']
        for col in required_columns:
            if col not in df.columns:
                print(f"错误: CSV文件缺少必要的列 '{col}'")
                return False
        
        # 连接数据库并创建表
        create_tables()
        
        # 导入数据
        print(f"开始导入数据，共 {len(df)} 条记录")
        records_imported = 0
        skipped = 0
        with db.atomic():
            # 批量插入，每批1000条记录
            batch_size = 1000
            for i in range(0, len(df), batch_size):
                batch = df.iloc[i:i+batch_size]
                data_to_insert = []
                
                for _, row in batch.iterrows():
                    # 检查是否已存在相同的交易哈希
                    try:
                        Transaction.get(Transaction.tx_hash == row['TxHash'])
                        skipped += 1
                        continue  # 如果存在，跳过此记录
                    except Transaction.DoesNotExist:
                        pass  # 如果不存在，继续处理
                    
                    # 准备数据
                    is_error = row.get('isError', 0) == 1 if 'isError' in row else False
                    
                    data_to_insert.append({
                        'tx_hash': row['TxHash'],
                        'block_height': row['BlockHeight'],
                        'timestamp': row['TimeStamp'],
                        'from_address': row['From'],
                        'to_address': row['To'],
                        'value': row['Value'],
                        'is_error': is_error
                    })
                
                if data_to_insert:
                    # 批量插入数据
                    Transaction.insert_many(data_to_insert).execute()
                    records_imported += len(data_to_insert)
                    print(f"已导入 {records_imported} 条记录")
        
        print(f"数据导入完成，共导入 {records_imported} 条记录")
        return {"imported":records_imported,"skipped":skipped}
    
    except Exception as e:
        print(f"导入数据时发生错误: {str(e)}")
        return False

def main():
    # 设置CSV文件路径
    csv_path = 'data/first_order_df.csv'
    
    # 导入数据
    success = import_data_from_csv(csv_path)
    
    if success:
        print("数据导入成功完成")
    else:
        print("数据导入失败")

if __name__ == "__main__":
    main()